from typing import List, Dict, Optional
import os
import glob
import ast



class FunctionNode:
    """Represents a function or method in Python code."""
    def __init__(self, name: str, module: str, class_name: Optional[str] = None):
        self.name = name
        self.module = module
        self.class_name = class_name
        self.full_name = f"{module}.{class_name}.{name}" if class_name else f"{module}.{name}"
        self.calls = set()  # Functions this function calls
        self.called_by = set()  # Functions that call this function
        self.source_code = ""
        self.lineno = 0
        self.end_lineno = 0

    def __repr__(self):
        return self.full_name

    def add_call(self, func_node: 'FunctionNode'):
        self.calls.add(func_node)
        func_node.called_by.add(self)

class CodeParser:
    """Parses Python code and builds function call relationships."""
    def __init__(self, src_directories: List[str], is_recursive: bool = True, local_only: bool = True):
        self.src_directories = src_directories
        self.is_recursive = is_recursive
        self.local_only = local_only
        self.functions = {}  # Maps function full_name to FunctionNode
        self.modules = {}  # Maps module name to file path
        self.classes = {}  # Maps class full_name to information about the class
        self.current_module = ""
        self.current_class = None
        self.file_content = ""
        self.entry_points = []  # To store identified entry points
        self.debug_mode = False  # Add debug mode flag

    def _debug_print(self, message):
        """Helper method for controlled debug output"""
        if self.debug_mode:
            print(message)

    def _is_local_file(self, file_path: str) -> bool:
        """Check if a file is within one of the source directories."""
        if not self.local_only:
            # If local_only is False, include all files
            return True
            
        file_path = os.path.abspath(file_path)
        for directory in self.src_directories:
            dir_path = os.path.abspath(directory)
            # Check if the file is within this directory
            if file_path.startswith(dir_path):
                return True
        return False

    def parse_all_files(self):
        """Parse all Python files in the specified directories."""
        # First pass - collect all functions
        for directory in self.src_directories:
            if not os.path.exists(directory):
                print(f"Warning: Directory {directory} does not exist")
                continue
            
            if self.is_recursive:
                pattern = os.path.join(directory, "**", "*.py")
                files = glob.glob(pattern, recursive=True)
            else:
                pattern = os.path.join(directory, "*.py")
                files = glob.glob(pattern)
            
            for file_path in files:
                # Only parse files that are local if local_only is True
                if self._is_local_file(file_path):
                    self.parse_file(file_path)
                elif self.debug_mode:
                    print(f"Skipping non-local file: {file_path}")
        
        # Print the functions we've collected for debugging
        self._debug_print(f"Collected functions: {list(self.functions.keys())}")
        
        # Second pass to resolve function calls
        for directory in self.src_directories:
            if not os.path.exists(directory):
                continue
            
            if self.is_recursive:
                pattern = os.path.join(directory, "**", "*.py")
                files = glob.glob(pattern, recursive=True)
            else:
                pattern = os.path.join(directory, "*.py")
                files = glob.glob(pattern)
            
            for file_path in files:
                # Only resolve calls in files that are local if local_only is True
                if self._is_local_file(file_path):
                    self.resolve_calls(file_path)
                elif self.debug_mode:
                    print(f"Skipping resolving calls in non-local file: {file_path}")
        
        # Turn on debug mode for better visibility
        self.debug_mode = False
        
        # Enhanced pass to resolve object method calls
        print("\nPerforming enhanced analysis of object method calls...")
        self._resolve_object_methods()
        
        if self.debug_mode:
            # Print a summary of the call graph
            self._print_call_graph_summary()

    def parse_file(self, file_path: str):
        """Parse a single Python file to extract functions and methods."""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                self.file_content = f.read()
            
            module_name = os.path.splitext(os.path.basename(file_path))[0]
            self.current_module = module_name
            self.modules[module_name] = file_path
            
            tree = ast.parse(self.file_content, filename=file_path)
            self.current_class = None
            
            # Ensure we capture calls in the main block
            if "__main__" in tree.body or "__main_block__" in tree.body:
                # Add special handling for main block
                main_block_name = f"{module_name}.__main_block__"
                self.functions[main_block_name] = {"calls": [], "source": ""}
                
                # Analyze the main block for function calls
                for node in tree.body:
                    if isinstance(node, ast.Expr) or isinstance(node, ast.Call) or isinstance(node, ast.Assign):
                        self._extract_calls_from_node(node, main_block_name)
            
            for node in ast.walk(tree):
                # Handle function definitions
                if isinstance(node, ast.FunctionDef):
                    func_node = self._process_function_def(node)
                    # Check if function might be an entry point (like 'main')
                    if func_node and func_node.name.lower() in ['main', 'run', 'start', 'execute']:
                        self.entry_points.append(func_node.full_name)
                
                # Handle class definitions
                elif isinstance(node, ast.ClassDef):
                    # Process the class itself first
                    class_info = self._process_class_def(node)
                    
                    # Then process its methods
                    self.current_class = node.name
                    for item in node.body:
                        if isinstance(item, ast.FunctionDef):
                            self._process_function_def(item)
                    self.current_class = None
        
        except Exception as e:
            print(f"Error parsing {file_path}: {e}")

    def _find_main_block(self, tree: ast.AST, file_path: str):
        """Identify and process __main__ block."""
        for node in ast.walk(tree):
            if isinstance(node, ast.If):
                # Check if this is an "if __name__ == '__main__':" block
                condition = None
                if isinstance(node.test, ast.Compare):
                    if isinstance(node.test.left, ast.Name) and node.test.left.id == '__name__':
                        if len(node.test.comparators) > 0 and isinstance(node.test.comparators[0], ast.Constant):
                            if node.test.comparators[0].value == '__main__':
                                condition = True
                
                if condition:
                    # Found a __main__ block, create a virtual function node for it
                    module_name = os.path.splitext(os.path.basename(file_path))[0]
                    main_block_name = f"{module_name}.__main_block__"
                    main_func = FunctionNode("__main_block__", module_name)
                    main_func.lineno = node.lineno
                    main_func.end_lineno = getattr(node, 'end_lineno', 0)
                    
                    # Get the source code for the main block
                    if main_func.end_lineno > 0:
                        source_lines = self.file_content.split('\n')[main_func.lineno-1:main_func.end_lineno]
                        main_func.source_code = '\n'.join(source_lines)
                    
                    self.functions[main_block_name] = main_func
                    self.entry_points.append(main_block_name)
                    return

    def _process_function_def(self, node: ast.FunctionDef):
        """Process a function definition node."""
        func_name = node.name
        class_name = self.current_class
        full_name = f"{self.current_module}.{class_name}.{func_name}" if class_name else f"{self.current_module}.{func_name}"
        
        # Create a new function node
        func_node = FunctionNode(func_name, self.current_module, class_name)
        func_node.lineno = node.lineno
        func_node.end_lineno = getattr(node, 'end_lineno', node.lineno)  # Python 3.8+ has end_lineno
        
        # Extract the source code
        source_lines = self.file_content.split('\n')[func_node.lineno-1:func_node.end_lineno]
        func_node.source_code = '\n'.join(source_lines)
        
        self.functions[full_name] = func_node
        return func_node

    def _process_class_def(self, node: ast.ClassDef):
        """Process a class definition node."""
        class_name = node.name
        full_name = f"{self.current_module}.{class_name}"
        
        # Store class information
        class_info = {
            'name': class_name,
            'module': self.current_module,
            'full_name': full_name,
            'methods': [],
            'lineno': node.lineno,
            'end_lineno': getattr(node, 'end_lineno', node.lineno),
            'source_code': ''
        }
        
        # Extract the source code
        source_lines = self.file_content.split('\n')[class_info['lineno']-1:class_info['end_lineno']]
        class_info['source_code'] = '\n'.join(source_lines)
        
        # Store class methods
        for item in node.body:
            if isinstance(item, ast.FunctionDef):
                method_name = item.name
                method_full_name = f"{full_name}.{method_name}"
                class_info['methods'].append(method_full_name)
        
        self.classes[full_name] = class_info
        
        self._debug_print(f"Processed class: {full_name} with methods: {class_info['methods']}")
        
        return class_info

    def resolve_calls(self, file_path: str):
        """Resolve function calls in a file."""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            module_name = os.path.splitext(os.path.basename(file_path))[0]
            self.current_module = module_name
            
            tree = ast.parse(content, filename=file_path)
            self.current_class = None
            
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    self.current_class = node.name
                    for item in node.body:
                        if isinstance(item, ast.FunctionDef):
                            self._resolve_calls_in_function(item)
                    self.current_class = None
                elif isinstance(node, ast.FunctionDef):
                    self._resolve_calls_in_function(node)
        
        except Exception as e:
            print(f"Error resolving calls in {file_path}: {e}")

    def _resolve_calls_in_function(self, func_node: ast.FunctionDef):
        """Resolve function calls within a function."""
        func_name = func_node.name
        class_name = self.current_class
        caller_full_name = f"{self.current_module}.{class_name}.{func_name}" if class_name else f"{self.current_module}.{func_name}"
        
        if caller_full_name not in self.functions:
            return
        
        caller = self.functions[caller_full_name]
        
        # Add import tracking dictionary
        imported_aliases = {}
        imported_modules = {}
        imported_classes = {}  # Track imported classes
        local_classes = {}     # Track locally defined classes
        
        # Track variable assignments to detect instantiations
        variable_types = {}
        
        # First, populate local_classes with classes defined in the current module
        for class_full_name, class_info in self.classes.items():
            if class_info['module'] == self.current_module:
                class_name = class_info['name']
                local_classes[class_name] = class_full_name
                self._debug_print(f"Found local class: {class_name} -> {class_full_name}")
        
        # Find all import statements at module level
        for node in ast.walk(ast.parse(self.file_content)):
            # Handle from X import Y cases
            if isinstance(node, ast.ImportFrom):
                module = node.module if node.module else ''
                
                # Handle wildcard imports (from X import *)
                if any(alias.name == '*' for alias in node.names):
                    self._debug_print(f"Found wildcard import: from {module} import *")
                    imported_names = self._process_wildcard_import(module)
                    
                    # Add all discovered names to the imported_aliases dictionary
                    for name in imported_names:
                        imported_aliases[name] = f"{module}.{name}"
                        # If name starts with uppercase, it's likely a class
                        if name and name[0].isupper():
                            imported_classes[name] = f"{module}.{name}"
                        # Ensure that the functions are added to the caller's calls
                        if name in self.functions:
                            caller.add_call(self.functions[f"{module}.{name}"])
                else:
                    for alias in node.names:
                        imported_name = alias.asname or alias.name
                        imported_aliases[imported_name] = f"{module}.{alias.name}" if module else alias.name
                        
                        # Track potential class imports
                        if imported_name[0].isupper():  # Classes typically start with uppercase
                            imported_classes[imported_name] = f"{module}.{alias.name}" if module else alias.name
            
            # Handle import X cases
            elif isinstance(node, ast.Import):
                for alias in node.names:
                    imported_name = alias.asname or alias.name
                    imported_modules[imported_name] = alias.name
        
        # First pass: track variable assignments to detect instantiations
        for node in ast.walk(func_node):
            if isinstance(node, ast.Assign):
                if isinstance(node.value, ast.Call):
                    # Handle class instantiation
                    if isinstance(node.value.func, ast.Name):
                        class_id = node.value.func.id
                        
                        # Check if it's a locally defined class
                        if class_id in local_classes:
                            for target in node.targets:
                                if isinstance(target, ast.Name):
                                    variable_types[target.id] = local_classes[class_id]
                                    self._debug_print(f"Detected variable {target.id} is instance of local class {local_classes[class_id]}")
                        
                        # Check if it's an imported class
                        elif class_id in imported_classes:
                            for target in node.targets:
                                if isinstance(target, ast.Name):
                                    variable_types[target.id] = imported_classes[class_id]
                                    self._debug_print(f"Detected variable {target.id} is instance of imported class {imported_classes[class_id]}")
        
        # Resolve calls
        for node in ast.walk(func_node):
            if isinstance(node, ast.Call):
                # Handle imported functions
                if isinstance(node.func, ast.Name):
                    called_name = node.func.id
                    
                    # Check if it's an imported alias (from X import Y)
                    if called_name in imported_aliases:
                        # Get the full module path and function name
                        imported_path = imported_aliases[called_name]
                        
                        # Try to resolve various possible formats for the callee
                        possible_callees = []
                        
                        # The direct mapping as specified in the import statement
                        possible_callees.append(imported_path)
                        
                        # If it's just the function name, try with the module name
                        if '.' not in imported_path:
                            parts = imported_path.split('.')
                            func_part = parts[-1]
                            for module_name in self.modules.keys():
                                if module_name.endswith(func_part) or module_name == func_part:
                                    possible_callees.append(f"{module_name}.{func_part}")
                        
                        # If local_only is True, filter callees to only include local modules
                        if self.local_only:
                            local_callees = []
                            for callee_name in possible_callees:
                                if '.' in callee_name:
                                    module = callee_name.split('.')[0]
                                    if module in self.modules:
                                        local_callees.append(callee_name)
                            possible_callees = local_callees if local_callees else possible_callees
                        
                        # Try to find the function in known functions
                        found = False
                        for callee_full_name in possible_callees:
                            if callee_full_name in self.functions:
                                caller.add_call(self.functions[callee_full_name])
                                found = True
                                break
                        
                        if not found and self.debug_mode:
                            print(f"Could not resolve imported function {called_name} -> {imported_path}")
                    
                    # Original resolution logic - check current module
                    else:
                        callee_full_name = f"{self.current_module}.{class_name}.{called_name}" if class_name else f"{self.current_module}.{called_name}"
                        if callee_full_name in self.functions:
                            caller.add_call(self.functions[callee_full_name])
                        else:
                            # Try to find the function across all modules
                            found = False
                            for func_full_name in self.functions.keys():
                                func_module, *rest = func_full_name.rsplit('.', 1)
                                func_short_name = rest[0] if rest else None
                                
                                if func_short_name == called_name:
                                    # If local_only is True, only include functions from local modules
                                    if self.local_only:
                                        # Extract the module name
                                        if '.' in func_full_name:
                                            module_name = func_full_name.split('.')[0]
                                            if module_name not in self.modules:
                                                continue
                                    
                                    # This could be the function we're looking for
                                    caller.add_call(self.functions[func_full_name])
                                    found = True
                                    break
                            
                            if not found and self.debug_mode:
                                print(f"Function {called_name} not found in any module")

                # Handle method and module function calls (e.g. obj.method(), module.func())
                elif isinstance(node.func, ast.Attribute):
                    # Handle method calls (e.g. self.method() or obj.method())
                    if isinstance(node.func.value, ast.Name):
                        obj_name = node.func.value.id
                        method_name = node.func.attr
                        
                        # Case 1: Method call within current class
                        if class_name and obj_name == "self":
                            callee_full_name = f"{self.current_module}.{class_name}.{method_name}"
                            if callee_full_name in self.functions:
                                caller.add_call(self.functions[callee_full_name])
                        
                        # Case 2: Method call on a variable with known type
                        elif obj_name in variable_types:
                            class_path = variable_types[obj_name]
                            method_full_name = f"{class_path}.{method_name}"
                            
                            self._debug_print(f"Looking for method: {method_full_name}")
                            
                            # Try direct match
                            if method_full_name in self.functions:
                                caller.add_call(self.functions[method_full_name])
                                self._debug_print(f"Connected call: {caller_full_name} -> {method_full_name}")
                            else:
                                # Try with just the class name (not full path)
                                class_name = class_path.split('.')[-1]
                                for known_func in self.functions.keys():
                                    parts = known_func.split('.')
                                    if len(parts) >= 2 and parts[-2] == class_name and parts[-1] == method_name:
                                        caller.add_call(self.functions[known_func])
                                        self._debug_print(f"Connected call via class name: {caller_full_name} -> {known_func}")
                                        break
                        else:
                            # Check if obj_name is a known imported module
                            if obj_name in imported_modules:
                                module_name = imported_modules[obj_name]
                                callee_full_name = f"{module_name}.{method_name}"
                                if callee_full_name in self.functions:
                                    caller.add_call(self.functions[callee_full_name])
                                else:
                                    # Try direct match
                                    callee_full_name = f"{obj_name}.{method_name}"  
                                    if callee_full_name in self.functions:
                                        caller.add_call(self.functions[callee_full_name])
                    
                    # Handle nested attribute calls (e.g. package.module.func())
                    elif isinstance(node.func.value, ast.Attribute):
                        # Try to resolve full dotted path
                        parts = []
                        current = node.func
                        while isinstance(current, ast.Attribute):
                            parts.append(current.attr)
                            current = current.value
                        if isinstance(current, ast.Name):
                            parts.append(current.id)
                        
                        # Build the full name from right to left
                        callee_full_name = ".".join(reversed(parts))
                        
                        # Try direct match
                        if callee_full_name in self.functions:
                            caller.add_call(self.functions[callee_full_name])
                        else:
                            # Try with the empty prefix only (no hardcoded prefixes)
                            prefixed_name = callee_full_name
                                
                            if prefixed_name in self.functions:
                                caller.add_call(self.functions[prefixed_name])

    def _process_wildcard_import(self, module_name: str) -> list:
        """Process a wildcard import by examining the module file to extract all available names."""
        self._debug_print(f"Processing wildcard import from module: {module_name}")
        imported_names = []
        
        # Try to locate the module file
        module_path = None
        module_parts = module_name.split('.')
        
        # First try to find it in known modules
        if module_name in self.modules:
            module_path = self.modules[module_name]
        else:
            # Try to look for the module in source directories
            for directory in self.src_directories:
                # Handle dot notation in module names
                potential_path = os.path.join(directory, *module_parts) + '.py'
                if os.path.exists(potential_path):
                    module_path = potential_path
                    break
                
                # Also try as a directory with __init__.py
                potential_init_path = os.path.join(directory, *module_parts, '__init__.py')
                if os.path.exists(potential_init_path):
                    module_path = potential_init_path
                    break
        
        if not module_path:
            self._debug_print(f"Could not locate module file for: {module_name}")
            return imported_names
        
        # Parse the module file to extract all top-level definitions
        try:
            with open(module_path, 'r', encoding='utf-8') as f:
                module_content = f.read()
                
            module_tree = ast.parse(module_content)
            
            # Extract all top-level function and class definitions
            for node in module_tree.body:
                if isinstance(node, ast.FunctionDef):
                    imported_names.append(node.name)
                    # Also store this function in our functions dictionary for later resolution
                    full_name = f"{module_name}.{node.name}"
                    if full_name not in self.functions:
                        func_node = FunctionNode(node.name, module_name)
                        func_node.lineno = node.lineno
                        func_node.end_lineno = getattr(node, 'end_lineno', node.lineno)
                        self.functions[full_name] = func_node
                
                elif isinstance(node, ast.ClassDef):
                    imported_names.append(node.name)
                    # Process methods within the class
                    for item in node.body:
                        if isinstance(item, ast.FunctionDef):
                            method_name = item.name
                            full_method_name = f"{module_name}.{node.name}.{method_name}"
                            if full_method_name not in self.functions:
                                method_node = FunctionNode(method_name, module_name, node.name)
                                method_node.lineno = item.lineno
                                method_node.end_lineno = getattr(item, 'end_lineno', item.lineno)
                                self.functions[full_method_name] = method_node
            
            self._debug_print(f"Extracted {len(imported_names)} names from {module_name}: {imported_names}")
        except Exception as e:
            self._debug_print(f"Error processing wildcard import from {module_name}: {e}")
        
        return imported_names

    def find_entry_point(self) -> str:
        """Find the most likely entry point if no root function is specified."""
        if not self.entry_points:
            # If no clear entry points, use the function with most incoming calls
            most_called = None
            max_callers = -1
            
            for func_name, func in self.functions.items():
                callers_count = len(func.called_by)
                if callers_count > max_callers:
                    max_callers = callers_count
                    most_called = func_name
            
            return most_called if most_called else next(iter(self.functions), None)
        
        # Prioritize __main__ blocks
        for entry in self.entry_points:
            if "__main_block__" in entry:
                return entry
        
        # Return the first entry point
        return self.entry_points[0]

    def print_function_catalog(self):
        """Print a catalog of all discovered functions and their calls for debugging."""
        print("\nDEBUG: Function Catalog")
        print("=====================")
        for func_name, func in sorted(self.functions.items()):
            print(f"Function: {func_name}")
            print(f"  Defined in: {func.module}{f'.{func.class_name}' if func.class_name else ''}")
            print(f"  Line: {func.lineno} - {func.end_lineno}")
            print(f"  Calls: {', '.join([f.full_name for f in func.calls]) if func.calls else 'None'}")
            print(f"  Called by: {', '.join([f.full_name for f in func.called_by]) if func.called_by else 'None'}")
            print("")

    def analyze_main_blocks(self):
        """Parse all Python files and analyze execution order in __main__ blocks."""
        # First ensure all files are parsed and calls are resolved
        self.parse_all_files()
        
        main_blocks = []
        execution_order = {}
        
        for directory in self.src_directories:
            if not os.path.exists(directory):
                print(f"Warning: Directory {directory} does not exist")
                continue
            
            if self.is_recursive:
                pattern = os.path.join(directory, "**", "*.py")
                files = glob.glob(pattern, recursive=True)
            else:
                pattern = os.path.join(directory, "*.py")
                files = glob.glob(pattern)
            
            for file_path in files:
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    module_name = os.path.splitext(os.path.basename(file_path))[0]
                    self.file_content = content
                    
                    tree = ast.parse(content, filename=file_path)
                    main_block = self._analyze_main_block(tree, file_path)
                    
                    if main_block:
                        main_blocks.append(main_block)
                        # Analyze execution order within this main block
                        execution_order[main_block] = self._determine_execution_order(main_block)
                
                except Exception as e:
                    print(f"Error analyzing main block in {file_path}: {e}")
        
        return main_blocks, execution_order

    def _analyze_main_block(self, tree: ast.AST, file_path: str) -> Optional[str]:
        """Identify and analyze __main__ block execution order."""
        for node in ast.walk(tree):
            if isinstance(node, ast.If):
                # Check if this is an "if __name__ == '__main__':" block
                if (isinstance(node.test, ast.Compare) and 
                    isinstance(node.test.left, ast.Name) and node.test.left.id == '__name__' and
                    len(node.test.comparators) > 0 and 
                    isinstance(node.test.comparators[0], ast.Constant) and 
                    node.test.comparators[0].value == '__main__'):
                    
                    # Found a __main__ block
                    module_name = os.path.splitext(os.path.basename(file_path))[0]
                    main_block_name = f"{module_name}.__main_block__"
                    
                    # If we haven't already processed this main block during parsing
                    if main_block_name not in self.functions:
                        main_func = FunctionNode("__main_block__", module_name)
                        main_func.lineno = node.lineno
                        main_func.end_lineno = getattr(node, 'end_lineno', 0)
                        
                        # Get the source code for the main block
                        if main_func.end_lineno > 0:
                            source_lines = self.file_content.split('\n')[main_func.lineno-1:main_func.end_lineno]
                            main_func.source_code = '\n'.join(source_lines)
                        
                        self.functions[main_block_name] = main_func
                        self.entry_points.append(main_block_name)
                    
                    # Analyze function calls within the main block
                    self._analyze_main_block_calls(node.body, main_block_name)
                    return main_block_name
        
        return None

    def _analyze_main_block_calls(self, body_nodes, main_block_name):
        """Extract function calls from the main block in execution order."""
        main_func = self.functions.get(main_block_name)
        if not main_func:
            return
        
        # Track imported functions for this module
        imported_aliases = {}
        imported_classes = {}  # Track imported classes
        imported_wildcard_modules = []  # Track modules imported with wildcard
        
        # Track variable assignments to detect object instantiations
        variable_types = {}
        
        # First identify imports at the module level
        module_name = main_block_name.split('.')[0]
        tree = ast.parse(self.file_content)
        
        # Print import analysis information
        print(f"\nAnalyzing imports in main block: {main_block_name}")
        
        # Extract all import statements
        for node in ast.walk(tree):
            if isinstance(node, ast.ImportFrom):
                module = node.module if node.module else ''
                
                # Special handling for wildcard imports
                if any(alias.name == '*' for alias in node.names):
                    imported_wildcard_modules.append(module)
                    print(f"Found import: from {module} import *")
                    
                    # Process wildcard import to extract functions
                    imported_names = self._process_wildcard_import(module)
                    for name in imported_names:
                        imported_aliases[name] = f"{module}.{name}"
                        # If name starts with uppercase, it's likely a class
                        if name and name[0].isupper():
                            imported_classes[name] = f"{module}.{name}"
                else:
                    for alias in node.names:
                        imported_name = alias.asname or alias.name
                        imported_aliases[imported_name] = f"{module}.{alias.name}" if module else alias.name
                        print(f"Found import: from {module} import {alias.name}{' as ' + alias.asname if alias.asname else ''}")
                        
                        # Track potential class imports
                        if imported_name[0].isupper():  # Classes typically start with uppercase
                            imported_classes[imported_name] = f"{module}.{alias.name}" if module else alias.name
            
            elif isinstance(node, ast.Import):
                for alias in node.names:
                    imported_name = alias.asname or alias.name
                    imported_aliases[imported_name] = alias.name
                    print(f"Found import: import {alias.name}{' as ' + alias.asname if alias.asname else ''}")
                    if '.' in alias.name:  # Handle submodule imports
                        parts = alias.name.split('.')
                        # Check for class imports in each submodule part
                        for part in parts:
                            if part[0].isupper():  # Potential class name
                                imported_classes[part] = alias.name
        
        # Add information about wildcard imports to debug output
        if imported_wildcard_modules:
            print(f"Wildcard imports from: {', '.join(imported_wildcard_modules)}")
        
        # Don't display * as an imported function
        if '*' in imported_aliases:
            del imported_aliases['*']
            
        print(f"All imported functions in {main_block_name}: {imported_aliases}")
        print(f"Imported classes: {imported_classes}")
        
        # First pass: track variable assignments to detect instantiations
        for node in body_nodes:
            if isinstance(node, ast.Assign):
                if isinstance(node.value, ast.Call):
                    # Handle class instantiation (e.g., analyzer = CallGraphAnalyzer(parser))
                    if isinstance(node.value.func, ast.Name) and node.value.func.id in imported_classes:
                        for target in node.targets:
                            if isinstance(target, ast.Name):
                                variable_types[target.id] = imported_classes[node.value.func.id]
                                print(f"DEBUG: Detected variable {target.id} is instance of {imported_classes[node.value.func.id]}")
                
                    # Handle instantiation via import (e.g., x = module.ClassName())
                    elif isinstance(node.value.func, ast.Attribute):
                        if isinstance(node.value.func.value, ast.Name):
                            module_part = node.value.func.value.id
                            class_part = node.value.func.attr
                            
                            if module_part in imported_aliases and class_part[0].isupper():
                                # This is likely a class instantiation
                                module_path = imported_aliases[module_part]
                                full_class_path = f"{module_path}.{class_part}"
                                
                                for target in node.targets:
                                    if isinstance(target, ast.Name):
                                        variable_types[target.id] = full_class_path
                                        print(f"DEBUG: Detected variable {target.id} is instance of {full_class_path}")
        
        # Check all statements in the main block
        print(f"Analyzing function calls in main block: {main_block_name}")
        for node in body_nodes:
            # Look for function calls in different statement types
            if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
                print(f"Found expression call: {ast.dump(node.value, annotate_fields=False)}")
                self._resolve_call_in_main_block(node.value, main_func, imported_aliases, variable_types)
            elif isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
                print(f"Found assignment call: {ast.dump(node.value, annotate_fields=False)}")
                self._resolve_call_in_main_block(node.value, main_func, imported_aliases, variable_types)
            elif isinstance(node, ast.If) or isinstance(node, ast.For) or isinstance(node, ast.While):
                # Recursively process control flow statements
                if isinstance(node, ast.If):
                    print(f"Processing if-statement body: {ast.dump(node.test, annotate_fields=False)}")
                self._analyze_main_block_calls(node.body, main_block_name)
                if hasattr(node, 'orelse') and node.orelse:
                    print(f"Processing else/elif body")
                    self._analyze_main_block_calls(node.orelse, main_block_name)
            elif isinstance(node, ast.With):
                # Handle with statements (e.g., with open() as f:)
                print(f"Processing with-statement body")
                self._analyze_main_block_calls(node.body, main_block_name)
                # Also check for function calls in the context expression
                for item in node.items:
                    if isinstance(item.context_expr, ast.Call):
                        print(f"Found with-context call: {ast.dump(item.context_expr, annotate_fields=False)}")
                        self._resolve_call_in_main_block(item.context_expr, main_func, imported_aliases, variable_types)

    def _resolve_call_in_main_block(self, call_node, main_func, imported_aliases=None, variable_types=None):
        """Resolve a function call in the main block and add it to the calls."""
        if imported_aliases is None:
            imported_aliases = {}
        if variable_types is None:
            variable_types = {}
        
        if isinstance(call_node.func, ast.Name):
            # Direct function call like main()
            called_name = call_node.func.id
            callee_full_name = f"{main_func.module}.{called_name}"
            
            # Output resolution attempt for debugging
            print(f"Resolving call to: {called_name}")
            
            # Check if this is an imported function
            if called_name in imported_aliases:
                imported_path = imported_aliases[called_name]
                
                # Try to find the function with various name patterns
                possible_callees = []
                
                # Direct mapping from import statement
                possible_callees.append(imported_path)
                
                # Try with module name variations
                parts = imported_path.split('.')
                if len(parts) > 1:
                    # For functions like 'Tools.print_info.print_configuration'
                    module_part = '.'.join(parts[:-1])
                    func_part = parts[-1]
                    possible_callees.append(f"{module_part}.{func_part}")
                    
                    # Try with just the module name and function
                    if len(parts) > 2:
                        possible_callees.append(f"{parts[-2]}.{parts[-1]}")
                        
                    # Try just the function name across all modules
                    for func_key in self.functions.keys():
                        if func_key.endswith(f".{func_part}"):
                            possible_callees.append(func_key)
                
                # Try to find the function in known functions
                found = False
                for callee_name in possible_callees:
                    if callee_name in self.functions:
                        main_func.add_call(self.functions[callee_name])
                        print(f"Found call: {called_name} -> {callee_name}")
                        found = True
                        break
                
                # If not found via imports, try more aggressively
                if not found:
                    for func_key in self.functions.keys():
                        # Match by just the function name at the end
                        if func_key.split('.')[-1] == called_name:
                            main_func.add_call(self.functions[func_key])
                            print(f"Found call by name match: {called_name} -> {func_key}")
                            found = True
                            break
            
            # If not imported, try current module
            elif callee_full_name in self.functions:
                main_func.add_call(self.functions[callee_full_name])
                print(f"Found local call: {called_name} -> {callee_full_name}")
            else:
                # Try to find the function across all modules
                found = False
                for func_full_name in self.functions.keys():
                    if func_full_name.endswith(f".{called_name}"):
                        main_func.add_call(self.functions[func_full_name])
                        print(f"Found call by suffix match: {called_name} -> {func_full_name}")
                        found = True
                        break
                        
                # For wildcard imports, try to find the function directly in the imported modules
                if not found:
                    for module_name in imported_aliases.values():
                        if '.' in module_name:
                            module_prefix = module_name.split('.')[0]
                            potential_name = f"{module_prefix}.{called_name}"
                            
                            if potential_name in self.functions:
                                main_func.add_call(self.functions[potential_name])
                                print(f"Found call in wildcard import: {called_name} -> {potential_name}")
                                found = True
                                break
                            
                            # Try deeper module structure
                            for func_key in self.functions.keys():
                                if func_key.startswith(f"{module_prefix}.") and func_key.endswith(f".{called_name}"):
                                    main_func.add_call(self.functions[func_key])
                                    print(f"Found call in nested wildcard import: {called_name} -> {func_key}")
                                    found = True
                                    break
        
        elif isinstance(call_node.func, ast.Attribute):
            # Method call like obj.method() or module.function()
            
            # Build the full call path
            call_parts = []
            current = call_node.func
            while isinstance(current, ast.Attribute):
                call_parts.append(current.attr)
                current = current.value
            
            if isinstance(current, ast.Name):
                obj_name = current.id
                method_name = call_parts[0]  # First part is the method name
                call_parts.append(obj_name)
                call_parts.reverse()
                
                # Handle method calls on tracked object instances
                if obj_name in variable_types:
                    class_path = variable_types[obj_name]
                    callee_full_name = f"{class_path}.{method_name}"
                    
                    print(f"DEBUG: Resolving method call: {obj_name}.{method_name} -> {callee_full_name}")
                    
                    if callee_full_name in self.functions:
                        main_func.add_call(self.functions[callee_full_name])
                        return
                    else:
                        # Try simplified class name
                        class_simple_name = class_path.split('.')[-1]
                        for func_key in self.functions:
                            parts = func_key.split('.')
                            if len(parts) > 2 and parts[-2] == class_simple_name and parts[-1] == method_name:
                                print(f"DEBUG: Found method via class name match: {func_key}")
                                main_func.add_call(self.functions[func_key])
                                return
                
                # Fall back to previous implementation if not a tracked object
                full_call_path = '.'.join(call_parts)
                
                # Try direct match first
                if full_call_path in self.functions:
                    main_func.add_call(self.functions[full_call_path])
                    return
                
                # Try with known module prefixes
                for prefix in ["Tools", "utils"]:
                    prefixed_call = f"{prefix}.{full_call_path}"
                    if prefixed_call in self.functions:
                        main_func.add_call(self.functions[prefixed_call])
                        return
                
                # Try matching just by the function name
                func_name = call_parts[-1]
                for func_key in self.functions.keys():
                    if func_key.endswith(f".{func_name}"):
                        # Check if the module part also matches
                        if len(call_parts) > 1:
                            module_part = call_parts[-2]
                            if module_part in func_key:
                                main_func.add_call(self.functions[func_key])
                                return
                        else:
                            # Just match by function name if no module specified
                            main_func.add_call(self.functions[func_key])
                            return

    def _determine_execution_order(self, main_block_name):
        """Determine the execution order of functions called from the main block."""
        main_func = self.functions.get(main_block_name)
        if not main_func:
            return []
        
        # First level is direct calls from main block
        execution_order = list(main_func.calls)
        
        # Add any nested calls to get a more complete picture
        all_reachable = set()
        queue = list(execution_order)
        visited = set()
        
        while queue:
            current = queue.pop(0)
            if current.full_name in visited:
                continue
            
            visited.add(current.full_name)
            all_reachable.add(current)
            
            for called in current.calls:
                if called.full_name not in visited:
                    queue.append(called)
        
        # Return full execution order
        return list(all_reachable)

    # Add a helper function to handle method chaining
    def _handle_method_chain(self, node, caller, variable_types=None):
        """Handle method chaining like obj.method1().method2()"""
        if variable_types is None:
            variable_types = {}
        
        # For a method chain, we need to process each call in the chain
        current_node = node
        chain = []
        
        # Extract the chain of method calls
        while isinstance(current_node, ast.Call) and isinstance(current_node.func, ast.Attribute):
            chain.append((current_node.func.attr, current_node.args, current_node.keywords))
            current_node = current_node.func.value
        
        # If we've found a chain ending in an object
        if isinstance(current_node, ast.Name) and current_node.id in variable_types:
            obj_name = current_node.id
            class_path = variable_types[obj_name]
            
            # Process each method in the chain (in reverse order)
            for method_name, args, keywords in reversed(chain):
                callee_full_name = f"{class_path}.{method_name}"
                
                if callee_full_name in self.functions:
                    caller.add_call(self.functions[callee_full_name])
                else:
                    # Try simplified class name
                    class_simple_name = class_path.split('.')[-1]
                    for func_key in self.functions:
                        parts = func_key.split('.')
                        if len(parts) > 2 and parts[-2] == class_simple_name and parts[-1] == method_name:
                            caller.add_call(self.functions[func_key])
                            break

    def _track_function_return_types(self):
        """Analyze functions to determine their return types for better method resolution."""
        return_types = {}
        
        for func_name, func_node in self.functions.items():
            # Skip if no source code
            if not func_node.source_code:
                continue
            
            # Parse the function to find return statements
            try:
                tree = ast.parse(func_node.source_code)
                
                # Check if this is a factory function that returns class instances
                for node in ast.walk(tree):
                    if isinstance(node, ast.Return) and isinstance(node.value, ast.Call):
                        if isinstance(node.value.func, ast.Name):
                            returned_class = node.value.func.id
                            # If the returned value looks like a class (starts with uppercase)
                            if returned_class[0].isupper():
                                # Record that this function returns instances of this class
                                return_types[func_name] = returned_class
            except:
                # Ignore parsing errors in return type analysis
                pass
            
        return return_types

    def _track_container_types(self, node, variable_types):
        """Track types of objects stored in lists, dictionaries, etc."""
        if isinstance(node, ast.Assign):
            # Handle list assignments with known types
            if isinstance(node.value, ast.List):
                # Check if list contains instances of tracked types
                potential_type = None
                for elt in node.value.elts:
                    if isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name):
                        class_name = elt.func.id
                        if class_name[0].isupper():  # Likely a class
                            potential_type = f"List[{class_name}]"
                            break
                            
                if potential_type:
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            variable_types[target.id] = potential_type
                            
            # Handle dictionary assignments
            elif isinstance(node.value, ast.Dict):
                # Check if this is a mapping of names to class instances
                value_type = None
                for key, val in zip(node.value.keys, node.value.values):
                    if isinstance(val, ast.Call) and isinstance(val.func, ast.Name):
                        class_name = val.func.id
                        if class_name[0].isupper():  # Likely a class
                            value_type = class_name
                            break
                            
                if value_type:
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            variable_types[target.id] = f"Dict[{value_type}]"

    def _resolve_object_methods(self):
        """Special pass to ensure object method calls are properly connected in the call graph."""
        # Focus especially on the analyze_python_code function
        target_functions = ["python_flow_map.analyze_python_code"]
        
        for func_name in target_functions:
            if func_name not in self.functions:
                print(f"Warning: Target function {func_name} not found")
                continue
            
            func_node = self.functions[func_name]
            
            # Get the source code
            source = func_node.source_code
            if not source:
                continue
            
            print(f"Analyzing object method calls in {func_name}")
            
            # Parse the source code
            try:
                tree = ast.parse(source)
                
                # Track variable assignments
                variable_types = {}
                
                # First pass: identify object instantiations (e.g., parser = CodeParser(...))
                for node in ast.walk(tree):
                    if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
                        if isinstance(node.value.func, ast.Name):
                            class_name = node.value.func.id
                            # Check if it's a class name
                            if class_name[0].isupper():
                                for target in node.targets:
                                    if isinstance(target, ast.Name):
                                        # Map variable to a potential class type
                                        var_name = target.id
                                        
                                        # First look for the class in our locally defined classes
                                        class_found = False
                                        for class_path in self.classes.keys():
                                            if class_path.endswith(f".{class_name}"):
                                                variable_types[var_name] = class_path
                                                print(f"Found locally defined class: {var_name} is of type {class_path}")
                                                class_found = True
                                                break
                                        
                                        # If not found in local classes, try to find it in known modules
                                        if not class_found:
                                            class_path = None
                                            for known_func in self.functions.keys():
                                                if class_name in known_func:
                                                    parts = known_func.split(".")
                                                    # If it looks like a class method (Class.method)
                                                    if len(parts) >= 2 and parts[-2] == class_name:
                                                        class_path = ".".join(parts[:-1])
                                                        break
                                            
                                            if class_path:
                                                variable_types[var_name] = class_path
                                                print(f"Detected variable {var_name} is of type {class_path}")
                                            else:
                                                # Use a basic mapping if exact path not found
                                                if class_name == "CodeParser":
                                                    variable_types[var_name] = "Tools.CodeParser"
                                                elif class_name == "CallGraphAnalyzer":
                                                    variable_types[var_name] = "Tools.CallGraphAnalyzer"
                                                elif class_name == "TreeVisualizer":
                                                    variable_types[var_name] = "Tools.TreeVisualizer"
                                                print(f"Mapped variable {var_name} to type {variable_types.get(var_name, 'unknown')}")
                
                # Second pass: connect method calls to their implementations
                for node in ast.walk(tree):
                    if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
                        if isinstance(node.func.value, ast.Name):
                            obj_name = node.func.value.id
                            method_name = node.func.attr
                            
                            if obj_name in variable_types:
                                class_path = variable_types[obj_name]
                                method_full_name = f"{class_path}.{method_name}"
                                
                                print(f"Looking for method: {method_full_name}")
                                
                                # Try direct match
                                if method_full_name in self.functions:
                                    func_node.add_call(self.functions[method_full_name])
                                    print(f"Connected call: {func_name} -> {method_full_name}")
                                else:
                                    # Try with just the class name (not full path)
                                    class_name = class_path.split('.')[-1]
                                    for known_func in self.functions.keys():
                                        parts = known_func.split('.')
                                        if len(parts) >= 2 and parts[-2] == class_name and parts[-1] == method_name:
                                            func_node.add_call(self.functions[known_func])
                                            print(f"Connected call via class name: {func_name} -> {known_func}")
                                            break
            except Exception as e:
                print(f"Error in enhanced analysis of {func_name}: {e}")

    def _print_call_graph_summary(self):
        """Print a summary of the call graph for debugging."""
        print("\nCall Graph Summary:")
        for func_name, func_node in sorted(self.functions.items()):
            if func_node.calls:
                print(f"{func_name} calls:")
                for called in func_node.calls:
                    print(f"  -> {called.full_name}")
            
            if func_node.called_by:
                print(f"{func_name} called by:")
                for caller in func_node.called_by:
                    print(f"  <- {caller.full_name}")

    def get_module_imports(self, module_name):
        """
        Returns a dictionary of imports for a given module.
        For wildcard imports, returns a list of imported modules under the '*' key.
        """
        imports = {}
        wildcard_modules = []
        
        # Try to get the file content for this module
        file_path = self.modules.get(module_name)
        if not file_path or not os.path.exists(file_path):
            return imports
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                file_content = f.read()
                
            tree = ast.parse(file_content)
            
            # Extract import statements
            for node in ast.walk(tree):
                if isinstance(node, ast.ImportFrom):
                    mod = node.module if node.module else ''
                    
                    # Handle wildcard imports
                    if any(alias.name == '*' for alias in node.names):
                        wildcard_modules.append(mod)
                    else:
                        for alias in node.names:
                            imports[alias.asname or alias.name] = f"{mod}.{alias.name}"
                
                elif isinstance(node, ast.Import):
                    for alias in node.names:
                        imports[alias.asname or alias.name] = alias.name
            
            # Add wildcard modules list if any were found
            if wildcard_modules:
                imports['*'] = wildcard_modules
                
        except Exception as e:
            print(f"Error extracting imports from {module_name}: {e}")
            
        return imports

    def process_wildcard_module_dependencies(self, calling_module, imported_module):
        """
        Process dependencies between a module and a wildcard-imported module.
        This ensures that functions called from wildcard imports are properly connected.
        """
        print(f"Processing dependencies: {calling_module} -> {imported_module}")
        
        # Get all functions from the calling module
        calling_funcs = [f for f in self.functions.keys() if f.startswith(f"{calling_module}.")]
        
        # Get all functions from the imported module
        imported_funcs = [f for f in self.functions.keys() if f.startswith(f"{imported_module}.")]
        
        if not calling_funcs or not imported_funcs:
            print(f"  No functions found in one of the modules.")
            return
            
        print(f"  Found {len(calling_funcs)} functions in {calling_module}")
        print(f"  Found {len(imported_funcs)} functions in {imported_module}")
        
        # For each function in the calling module
        for caller_name in calling_funcs:
            caller = self.functions.get(caller_name)
            if not caller:
                continue
                
            # Get the source code
            if not caller.source_code:
                continue
                
            # For each function in the imported module
            for imported_func in imported_funcs:
                imported_name = imported_func.split('.')[-1]  # Get just the function name
                
                # Check if the caller's source contains direct calls to this imported function
                # This is a simple string-based check that can be improved
                if f"{imported_name}(" in caller.source_code:
                    imported_node = self.functions.get(imported_func)
                    if imported_node:
                        caller.add_call(imported_node)
                        print(f"  Connected: {caller_name} -> {imported_func}")
                
                # Also check for common patterns like passing the function as argument
                if f"{imported_name}," in caller.source_code or f"{imported_name})" in caller.source_code:
                    imported_node = self.functions.get(imported_func)
                    if imported_node:
                        caller.add_call(imported_node)
                        print(f"  Connected (as arg): {caller_name} -> {imported_func}")

